{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d0c6d29b",
   "metadata": {},
   "source": [
    "# CSWin Transformer 创新点与影响总结\n",
    "\n",
    "## 主要创新点\n",
    "\n",
    "  ![alt text](resources/cswin_arch.png \"Title\")\n",
    "\n",
    "\n",
    "### 1. **十字形窗口自注意力机制（Cross-Shaped Window Attention）**\n",
    "- 核心思想：通过水平和垂直条带的并行处理，构建交叉形状的注意力窗口，显著扩大每个token的注意力范围，同时降低计算复杂度。\n",
    "- 优势：相比传统全局注意力机制，计算效率更高；相比局部注意力机制，能缓解感受野受限的问题，提升长距离依赖建模能力。\n",
    "\n",
    "  ![alt text](resources/cswin_attention_window.png \"Title\")\n",
    "\n",
    "### 2. **局部增强位置编码（Locally Enhanced Positional Encoding, LePE）**\n",
    "- 设计目的：增强局部区域的位置信息建模，弥补交叉形窗口可能忽略的局部细节。\n",
    "- 实现方式：在自注意力计算后引入局部位置信息，通过卷积操作提取局部特征并融合到输出中。\n",
    "\n",
    "\n",
    "  ![alt text](resources/cswin_lepe.png \"Title\")\n",
    "\n",
    "### 3. **多尺度条带宽度设计**\n",
    "- 理论支持：通过数学分析条带宽度对建模能力的影响，提出不同网络层动态调整条带宽度的策略，进一步平衡计算成本与性能。\n",
    "\n",
    "---\n",
    "\n",
    "## 模型影响与意义\n",
    "\n",
    "### 1. **计算效率与性能的平衡**\n",
    "- CSWin Transformer 在视觉任务中实现了高效的全局-局部交互，在ImageNet分类、COCO检测等任务中达到SOTA性能，同时保持较低的计算开销。\n",
    "\n",
    "### 2. **推动Transformer架构改进**\n",
    "- 为后续研究提供新思路，例如在YOLOv8等检测框架中引入CSWin模块，提升主干网络的多尺度特征提取能力。\n",
    "\n",
    "### 3. **开源与应用价值**\n",
    "- 作为CVPR 2022收录论文，其代码和方法已被广泛复现，应用于人脸修复、医学图像分析等领域，验证了其通用性。\n",
    "\n",
    "> **总结**：CSWin Transformer 通过创新性的交叉形窗口设计和位置编码策略，解决了视觉Transformer中全局注意力计算量大与局部注意力感受野受限的矛盾，为高效视觉模型设计提供了重要参考。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "12d71aad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n",
      "Use device:  cuda\n"
     ]
    }
   ],
   "source": [
    "# 自动重新加载外部module，使得修改代码之后无需重新import\n",
    "# see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from hdd.device.utils import get_device\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# 设置训练数据的路径\n",
    "DATA_ROOT = \"~/workspace/hands-dirty-on-dl/dataset\"\n",
    "# 设置TensorBoard的路径\n",
    "TENSORBOARD_ROOT = \"~/workspace/hands-dirty-on-dl/dataset\"\n",
    "# 设置预训练模型参数路径\n",
    "TORCH_HUB_PATH = \"~/workspace/hands-dirty-on-dl/pretrained_models\"\n",
    "torch.hub.set_dir(TORCH_HUB_PATH)\n",
    "# 挑选最合适的训练设备\n",
    "DEVICE = get_device([\"cuda\", \"cpu\"])\n",
    "max_epochs = 150\n",
    "print(\"Use device: \", DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "42716cbb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "from hdd.data_util.auto_augmentation import CIFAR10Policy\n",
    "\n",
    "# 训练超参数和数据增强来自 https://github.com/omihub777/ViT-CIFAR\n",
    "CIFAR_10_MEAN = [0.4914, 0.4822, 0.4465]\n",
    "CIFAR_10_STD = [0.2470, 0.2435, 0.2616]\n",
    "BATCH_SIZE = 128\n",
    "\n",
    "val_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(CIFAR_10_MEAN, CIFAR_10_STD),\n",
    "    ]\n",
    ")\n",
    "\n",
    "val_dataloader = torch.utils.data.DataLoader(\n",
    "    datasets.CIFAR10(\n",
    "        root=DATA_ROOT, train=False, download=True, transform=val_transform\n",
    "    ),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=False,\n",
    "    num_workers=8,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "train_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.RandomCrop(size=32, padding=4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        CIFAR10Policy(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(CIFAR_10_MEAN, CIFAR_10_STD),\n",
    "    ]\n",
    ")\n",
    "\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    datasets.CIFAR10(\n",
    "        root=DATA_ROOT, train=True, download=True, transform=train_transform\n",
    "    ),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    num_workers=8,\n",
    "    pin_memory=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ce034d25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#Parameter: 6488218\n",
      "Epoch: 1/150 Train Loss: 2.3023 Accuracy: 0.1209 Time: 16.95441  | Val Loss: 2.2906 Accuracy: 0.1328\n",
      "Epoch: 2/150 Train Loss: 2.0738 Accuracy: 0.2553 Time: 17.46283  | Val Loss: 1.7979 Accuracy: 0.3928\n",
      "Epoch: 3/150 Train Loss: 1.8749 Accuracy: 0.3539 Time: 16.72593  | Val Loss: 1.6264 Accuracy: 0.4791\n",
      "Epoch: 4/150 Train Loss: 1.7130 Accuracy: 0.4390 Time: 16.38400  | Val Loss: 1.4757 Accuracy: 0.5578\n",
      "Epoch: 5/150 Train Loss: 1.5824 Accuracy: 0.5038 Time: 16.71232  | Val Loss: 1.3418 Accuracy: 0.6171\n",
      "Epoch: 6/150 Train Loss: 1.5204 Accuracy: 0.5358 Time: 16.77067  | Val Loss: 1.3044 Accuracy: 0.6341\n",
      "Epoch: 7/150 Train Loss: 1.4693 Accuracy: 0.5607 Time: 16.44648  | Val Loss: 1.2820 Accuracy: 0.6494\n",
      "Epoch: 8/150 Train Loss: 1.4364 Accuracy: 0.5757 Time: 16.24806  | Val Loss: 1.2818 Accuracy: 0.6394\n",
      "Epoch: 9/150 Train Loss: 1.4046 Accuracy: 0.5925 Time: 16.83452  | Val Loss: 1.2033 Accuracy: 0.6844\n",
      "Epoch: 10/150 Train Loss: 1.3715 Accuracy: 0.6058 Time: 16.21724  | Val Loss: 1.2081 Accuracy: 0.6820\n",
      "Epoch: 11/150 Train Loss: 1.3452 Accuracy: 0.6203 Time: 16.41993  | Val Loss: 1.1746 Accuracy: 0.6973\n",
      "Epoch: 12/150 Train Loss: 1.2882 Accuracy: 0.6470 Time: 16.37248  | Val Loss: 1.1015 Accuracy: 0.7374\n",
      "Epoch: 13/150 Train Loss: 1.2547 Accuracy: 0.6623 Time: 16.77938  | Val Loss: 1.0790 Accuracy: 0.7426\n",
      "Epoch: 14/150 Train Loss: 1.2123 Accuracy: 0.6817 Time: 17.51064  | Val Loss: 1.0425 Accuracy: 0.7635\n",
      "Epoch: 15/150 Train Loss: 1.1747 Accuracy: 0.6971 Time: 16.67213  | Val Loss: 1.0076 Accuracy: 0.7788\n",
      "Epoch: 16/150 Train Loss: 1.1495 Accuracy: 0.7106 Time: 17.01767  | Val Loss: 0.9733 Accuracy: 0.7941\n",
      "Epoch: 17/150 Train Loss: 1.1242 Accuracy: 0.7213 Time: 17.08502  | Val Loss: 0.9545 Accuracy: 0.8013\n",
      "Epoch: 18/150 Train Loss: 1.0945 Accuracy: 0.7353 Time: 16.66516  | Val Loss: 0.9624 Accuracy: 0.7986\n",
      "Epoch: 19/150 Train Loss: 1.0685 Accuracy: 0.7484 Time: 16.28951  | Val Loss: 0.9246 Accuracy: 0.8151\n",
      "Epoch: 20/150 Train Loss: 1.0455 Accuracy: 0.7572 Time: 16.68889  | Val Loss: 0.9458 Accuracy: 0.8067\n",
      "Epoch: 21/150 Train Loss: 1.0295 Accuracy: 0.7645 Time: 15.66224  | Val Loss: 0.8970 Accuracy: 0.8254\n",
      "Epoch: 22/150 Train Loss: 1.0106 Accuracy: 0.7714 Time: 15.73402  | Val Loss: 0.8761 Accuracy: 0.8358\n",
      "Epoch: 23/150 Train Loss: 0.9949 Accuracy: 0.7792 Time: 16.10529  | Val Loss: 0.9143 Accuracy: 0.8173\n",
      "Epoch: 24/150 Train Loss: 0.9786 Accuracy: 0.7881 Time: 15.86618  | Val Loss: 0.8500 Accuracy: 0.8423\n",
      "Epoch: 25/150 Train Loss: 0.9701 Accuracy: 0.7913 Time: 15.78540  | Val Loss: 0.8429 Accuracy: 0.8497\n",
      "Epoch: 26/150 Train Loss: 0.9550 Accuracy: 0.7972 Time: 15.84344  | Val Loss: 0.8382 Accuracy: 0.8498\n",
      "Epoch: 27/150 Train Loss: 0.9423 Accuracy: 0.8026 Time: 15.93502  | Val Loss: 0.8394 Accuracy: 0.8458\n",
      "Epoch: 28/150 Train Loss: 0.9355 Accuracy: 0.8067 Time: 15.47917  | Val Loss: 0.8338 Accuracy: 0.8533\n",
      "Epoch: 29/150 Train Loss: 0.9171 Accuracy: 0.8144 Time: 16.21065  | Val Loss: 0.8246 Accuracy: 0.8545\n",
      "Epoch: 30/150 Train Loss: 0.9092 Accuracy: 0.8178 Time: 15.96844  | Val Loss: 0.8159 Accuracy: 0.8612\n",
      "Epoch: 31/150 Train Loss: 0.9009 Accuracy: 0.8224 Time: 15.54464  | Val Loss: 0.8014 Accuracy: 0.8644\n",
      "Epoch: 32/150 Train Loss: 0.8951 Accuracy: 0.8251 Time: 15.46677  | Val Loss: 0.8177 Accuracy: 0.8573\n",
      "Epoch: 33/150 Train Loss: 0.8794 Accuracy: 0.8310 Time: 16.19838  | Val Loss: 0.8020 Accuracy: 0.8696\n",
      "Epoch: 34/150 Train Loss: 0.8731 Accuracy: 0.8352 Time: 15.93771  | Val Loss: 0.7776 Accuracy: 0.8793\n",
      "Epoch: 35/150 Train Loss: 0.8687 Accuracy: 0.8359 Time: 15.82116  | Val Loss: 0.7882 Accuracy: 0.8758\n",
      "Epoch: 36/150 Train Loss: 0.8562 Accuracy: 0.8416 Time: 15.83626  | Val Loss: 0.7608 Accuracy: 0.8852\n",
      "Epoch: 37/150 Train Loss: 0.8521 Accuracy: 0.8425 Time: 16.31287  | Val Loss: 0.7765 Accuracy: 0.8827\n",
      "Epoch: 38/150 Train Loss: 0.8368 Accuracy: 0.8513 Time: 15.94816  | Val Loss: 0.7502 Accuracy: 0.8890\n",
      "Epoch: 39/150 Train Loss: 0.8342 Accuracy: 0.8519 Time: 15.58628  | Val Loss: 0.7718 Accuracy: 0.8770\n",
      "Epoch: 40/150 Train Loss: 0.8258 Accuracy: 0.8560 Time: 16.30549  | Val Loss: 0.7661 Accuracy: 0.8811\n",
      "Epoch: 41/150 Train Loss: 0.8223 Accuracy: 0.8582 Time: 15.74261  | Val Loss: 0.7651 Accuracy: 0.8833\n",
      "Epoch: 42/150 Train Loss: 0.8126 Accuracy: 0.8604 Time: 15.96690  | Val Loss: 0.7478 Accuracy: 0.8933\n",
      "Epoch: 43/150 Train Loss: 0.8063 Accuracy: 0.8637 Time: 15.77491  | Val Loss: 0.7407 Accuracy: 0.8940\n",
      "Epoch: 44/150 Train Loss: 0.8013 Accuracy: 0.8655 Time: 16.09503  | Val Loss: 0.7530 Accuracy: 0.8900\n",
      "Epoch: 45/150 Train Loss: 0.7974 Accuracy: 0.8689 Time: 15.94421  | Val Loss: 0.7424 Accuracy: 0.8943\n",
      "Epoch: 46/150 Train Loss: 0.7905 Accuracy: 0.8709 Time: 15.92503  | Val Loss: 0.7529 Accuracy: 0.8923\n",
      "Epoch: 47/150 Train Loss: 0.7907 Accuracy: 0.8703 Time: 15.96442  | Val Loss: 0.7255 Accuracy: 0.9007\n",
      "Epoch: 48/150 Train Loss: 0.7792 Accuracy: 0.8772 Time: 15.55521  | Val Loss: 0.7302 Accuracy: 0.8985\n",
      "Epoch: 49/150 Train Loss: 0.7763 Accuracy: 0.8765 Time: 15.50747  | Val Loss: 0.7382 Accuracy: 0.8963\n",
      "Epoch: 50/150 Train Loss: 0.7697 Accuracy: 0.8807 Time: 16.22556  | Val Loss: 0.7356 Accuracy: 0.8984\n",
      "Epoch: 51/150 Train Loss: 0.7685 Accuracy: 0.8805 Time: 16.43595  | Val Loss: 0.7154 Accuracy: 0.9059\n",
      "Epoch: 52/150 Train Loss: 0.7607 Accuracy: 0.8830 Time: 18.06986  | Val Loss: 0.7316 Accuracy: 0.8999\n",
      "Epoch: 53/150 Train Loss: 0.7561 Accuracy: 0.8857 Time: 16.36459  | Val Loss: 0.7295 Accuracy: 0.9027\n",
      "Epoch: 54/150 Train Loss: 0.7523 Accuracy: 0.8875 Time: 16.85250  | Val Loss: 0.7141 Accuracy: 0.9073\n",
      "Epoch: 55/150 Train Loss: 0.7505 Accuracy: 0.8885 Time: 16.42504  | Val Loss: 0.7264 Accuracy: 0.9026\n",
      "Epoch: 56/150 Train Loss: 0.7436 Accuracy: 0.8914 Time: 17.54026  | Val Loss: 0.7128 Accuracy: 0.9090\n",
      "Epoch: 57/150 Train Loss: 0.7380 Accuracy: 0.8936 Time: 16.83167  | Val Loss: 0.7058 Accuracy: 0.9125\n",
      "Epoch: 58/150 Train Loss: 0.7326 Accuracy: 0.8973 Time: 17.41557  | Val Loss: 0.7163 Accuracy: 0.9094\n",
      "Epoch: 59/150 Train Loss: 0.7320 Accuracy: 0.8962 Time: 17.03982  | Val Loss: 0.7296 Accuracy: 0.8987\n",
      "Epoch: 60/150 Train Loss: 0.7265 Accuracy: 0.8995 Time: 16.84844  | Val Loss: 0.7178 Accuracy: 0.9074\n",
      "Epoch: 61/150 Train Loss: 0.7240 Accuracy: 0.9006 Time: 17.40874  | Val Loss: 0.7100 Accuracy: 0.9130\n",
      "Epoch: 62/150 Train Loss: 0.7200 Accuracy: 0.9033 Time: 17.12774  | Val Loss: 0.7006 Accuracy: 0.9161\n",
      "Epoch: 63/150 Train Loss: 0.7164 Accuracy: 0.9038 Time: 16.87647  | Val Loss: 0.7184 Accuracy: 0.9073\n",
      "Epoch: 64/150 Train Loss: 0.7122 Accuracy: 0.9058 Time: 17.46002  | Val Loss: 0.7122 Accuracy: 0.9092\n",
      "Epoch: 65/150 Train Loss: 0.7098 Accuracy: 0.9065 Time: 18.17431  | Val Loss: 0.7016 Accuracy: 0.9163\n",
      "Epoch: 66/150 Train Loss: 0.7061 Accuracy: 0.9080 Time: 18.16259  | Val Loss: 0.6976 Accuracy: 0.9176\n",
      "Epoch: 67/150 Train Loss: 0.7026 Accuracy: 0.9098 Time: 17.37134  | Val Loss: 0.7000 Accuracy: 0.9131\n",
      "Epoch: 68/150 Train Loss: 0.6982 Accuracy: 0.9119 Time: 17.41864  | Val Loss: 0.6964 Accuracy: 0.9199\n",
      "Epoch: 69/150 Train Loss: 0.6968 Accuracy: 0.9127 Time: 16.92233  | Val Loss: 0.6890 Accuracy: 0.9177\n",
      "Epoch: 70/150 Train Loss: 0.6955 Accuracy: 0.9133 Time: 18.26757  | Val Loss: 0.6899 Accuracy: 0.9183\n",
      "Epoch: 71/150 Train Loss: 0.6909 Accuracy: 0.9160 Time: 16.77824  | Val Loss: 0.6925 Accuracy: 0.9188\n",
      "Epoch: 72/150 Train Loss: 0.6890 Accuracy: 0.9153 Time: 16.65785  | Val Loss: 0.7020 Accuracy: 0.9146\n",
      "Epoch: 73/150 Train Loss: 0.6857 Accuracy: 0.9181 Time: 16.74914  | Val Loss: 0.6933 Accuracy: 0.9185\n",
      "Epoch: 74/150 Train Loss: 0.6811 Accuracy: 0.9197 Time: 16.63941  | Val Loss: 0.6850 Accuracy: 0.9218\n",
      "Epoch: 75/150 Train Loss: 0.6832 Accuracy: 0.9179 Time: 16.32077  | Val Loss: 0.6926 Accuracy: 0.9200\n",
      "Epoch: 76/150 Train Loss: 0.6725 Accuracy: 0.9238 Time: 16.92750  | Val Loss: 0.6903 Accuracy: 0.9198\n",
      "Epoch: 77/150 Train Loss: 0.6733 Accuracy: 0.9231 Time: 16.68258  | Val Loss: 0.7100 Accuracy: 0.9149\n",
      "Epoch: 78/150 Train Loss: 0.6696 Accuracy: 0.9245 Time: 16.98850  | Val Loss: 0.6992 Accuracy: 0.9170\n",
      "Epoch: 79/150 Train Loss: 0.6653 Accuracy: 0.9256 Time: 17.34134  | Val Loss: 0.7042 Accuracy: 0.9170\n",
      "Epoch: 80/150 Train Loss: 0.6661 Accuracy: 0.9261 Time: 17.74337  | Val Loss: 0.6958 Accuracy: 0.9182\n",
      "Epoch: 81/150 Train Loss: 0.6653 Accuracy: 0.9269 Time: 16.95755  | Val Loss: 0.6862 Accuracy: 0.9239\n",
      "Epoch: 82/150 Train Loss: 0.6597 Accuracy: 0.9300 Time: 16.97294  | Val Loss: 0.6992 Accuracy: 0.9160\n",
      "Epoch: 83/150 Train Loss: 0.6579 Accuracy: 0.9310 Time: 17.35184  | Val Loss: 0.6884 Accuracy: 0.9248\n",
      "Epoch: 84/150 Train Loss: 0.6543 Accuracy: 0.9324 Time: 16.93856  | Val Loss: 0.6986 Accuracy: 0.9220\n",
      "Epoch: 85/150 Train Loss: 0.6549 Accuracy: 0.9321 Time: 16.48703  | Val Loss: 0.6933 Accuracy: 0.9205\n",
      "Epoch: 86/150 Train Loss: 0.6508 Accuracy: 0.9329 Time: 17.39171  | Val Loss: 0.6843 Accuracy: 0.9256\n",
      "Epoch: 87/150 Train Loss: 0.6535 Accuracy: 0.9317 Time: 17.31792  | Val Loss: 0.6766 Accuracy: 0.9281\n",
      "Epoch: 88/150 Train Loss: 0.6476 Accuracy: 0.9342 Time: 16.83195  | Val Loss: 0.6916 Accuracy: 0.9209\n",
      "Epoch: 89/150 Train Loss: 0.6412 Accuracy: 0.9377 Time: 16.72709  | Val Loss: 0.6810 Accuracy: 0.9264\n",
      "Epoch: 90/150 Train Loss: 0.6411 Accuracy: 0.9361 Time: 16.87745  | Val Loss: 0.6800 Accuracy: 0.9251\n",
      "Epoch: 91/150 Train Loss: 0.6378 Accuracy: 0.9382 Time: 16.80333  | Val Loss: 0.6859 Accuracy: 0.9268\n",
      "Epoch: 92/150 Train Loss: 0.6379 Accuracy: 0.9391 Time: 17.00136  | Val Loss: 0.6786 Accuracy: 0.9265\n",
      "Epoch: 93/150 Train Loss: 0.6322 Accuracy: 0.9416 Time: 17.17902  | Val Loss: 0.6973 Accuracy: 0.9217\n",
      "Epoch: 94/150 Train Loss: 0.6319 Accuracy: 0.9415 Time: 17.61739  | Val Loss: 0.6761 Accuracy: 0.9296\n",
      "Epoch: 95/150 Train Loss: 0.6294 Accuracy: 0.9420 Time: 17.09540  | Val Loss: 0.6853 Accuracy: 0.9265\n",
      "Epoch: 96/150 Train Loss: 0.6301 Accuracy: 0.9429 Time: 17.12266  | Val Loss: 0.6854 Accuracy: 0.9246\n",
      "Epoch: 97/150 Train Loss: 0.6264 Accuracy: 0.9440 Time: 17.32571  | Val Loss: 0.6839 Accuracy: 0.9279\n",
      "Epoch: 98/150 Train Loss: 0.6232 Accuracy: 0.9456 Time: 17.49170  | Val Loss: 0.6821 Accuracy: 0.9283\n",
      "Epoch: 99/150 Train Loss: 0.6244 Accuracy: 0.9448 Time: 17.26189  | Val Loss: 0.6866 Accuracy: 0.9273\n",
      "Epoch: 100/150 Train Loss: 0.6225 Accuracy: 0.9458 Time: 16.98629  | Val Loss: 0.6843 Accuracy: 0.9249\n",
      "Epoch: 101/150 Train Loss: 0.6200 Accuracy: 0.9473 Time: 16.84403  | Val Loss: 0.6914 Accuracy: 0.9257\n",
      "Epoch: 102/150 Train Loss: 0.6148 Accuracy: 0.9491 Time: 17.40612  | Val Loss: 0.6823 Accuracy: 0.9299\n",
      "Epoch: 103/150 Train Loss: 0.6136 Accuracy: 0.9495 Time: 17.28136  | Val Loss: 0.6802 Accuracy: 0.9285\n",
      "Epoch: 104/150 Train Loss: 0.6132 Accuracy: 0.9506 Time: 17.22988  | Val Loss: 0.6847 Accuracy: 0.9269\n",
      "Epoch: 105/150 Train Loss: 0.6126 Accuracy: 0.9510 Time: 17.44362  | Val Loss: 0.6777 Accuracy: 0.9319\n",
      "Epoch: 106/150 Train Loss: 0.6118 Accuracy: 0.9507 Time: 17.37972  | Val Loss: 0.6808 Accuracy: 0.9293\n",
      "Epoch: 107/150 Train Loss: 0.6071 Accuracy: 0.9529 Time: 16.57580  | Val Loss: 0.6768 Accuracy: 0.9325\n",
      "Epoch: 108/150 Train Loss: 0.6094 Accuracy: 0.9514 Time: 16.94428  | Val Loss: 0.6698 Accuracy: 0.9352\n",
      "Epoch: 109/150 Train Loss: 0.6059 Accuracy: 0.9530 Time: 16.58662  | Val Loss: 0.6742 Accuracy: 0.9330\n",
      "Epoch: 110/150 Train Loss: 0.6034 Accuracy: 0.9535 Time: 16.52304  | Val Loss: 0.6714 Accuracy: 0.9321\n",
      "Epoch: 111/150 Train Loss: 0.6045 Accuracy: 0.9536 Time: 16.44442  | Val Loss: 0.6762 Accuracy: 0.9311\n",
      "Epoch: 112/150 Train Loss: 0.6015 Accuracy: 0.9551 Time: 16.37672  | Val Loss: 0.6755 Accuracy: 0.9319\n",
      "Epoch: 113/150 Train Loss: 0.6004 Accuracy: 0.9554 Time: 16.54423  | Val Loss: 0.6802 Accuracy: 0.9301\n",
      "Epoch: 114/150 Train Loss: 0.6003 Accuracy: 0.9558 Time: 16.42849  | Val Loss: 0.6682 Accuracy: 0.9355\n",
      "Epoch: 115/150 Train Loss: 0.5969 Accuracy: 0.9564 Time: 16.41514  | Val Loss: 0.6681 Accuracy: 0.9357\n",
      "Epoch: 116/150 Train Loss: 0.5951 Accuracy: 0.9577 Time: 16.75645  | Val Loss: 0.6675 Accuracy: 0.9357\n",
      "Epoch: 117/150 Train Loss: 0.5978 Accuracy: 0.9567 Time: 16.40054  | Val Loss: 0.6760 Accuracy: 0.9301\n",
      "Epoch: 118/150 Train Loss: 0.5994 Accuracy: 0.9555 Time: 16.77999  | Val Loss: 0.6677 Accuracy: 0.9365\n",
      "Epoch: 119/150 Train Loss: 0.5934 Accuracy: 0.9583 Time: 16.51627  | Val Loss: 0.6659 Accuracy: 0.9364\n",
      "Epoch: 120/150 Train Loss: 0.5913 Accuracy: 0.9596 Time: 17.13392  | Val Loss: 0.6672 Accuracy: 0.9346\n",
      "Epoch: 121/150 Train Loss: 0.5911 Accuracy: 0.9601 Time: 17.19112  | Val Loss: 0.6668 Accuracy: 0.9353\n",
      "Epoch: 122/150 Train Loss: 0.5882 Accuracy: 0.9605 Time: 16.85050  | Val Loss: 0.6705 Accuracy: 0.9335\n",
      "Epoch: 123/150 Train Loss: 0.5877 Accuracy: 0.9610 Time: 16.51614  | Val Loss: 0.6668 Accuracy: 0.9380\n",
      "Epoch: 124/150 Train Loss: 0.5879 Accuracy: 0.9607 Time: 16.84744  | Val Loss: 0.6725 Accuracy: 0.9342\n",
      "Epoch: 125/150 Train Loss: 0.5840 Accuracy: 0.9634 Time: 16.76942  | Val Loss: 0.6745 Accuracy: 0.9340\n",
      "Epoch: 126/150 Train Loss: 0.5855 Accuracy: 0.9622 Time: 16.76454  | Val Loss: 0.6711 Accuracy: 0.9340\n",
      "Epoch: 127/150 Train Loss: 0.5824 Accuracy: 0.9639 Time: 16.37874  | Val Loss: 0.6701 Accuracy: 0.9364\n",
      "Epoch: 128/150 Train Loss: 0.5839 Accuracy: 0.9624 Time: 17.07471  | Val Loss: 0.6717 Accuracy: 0.9358\n",
      "Epoch: 129/150 Train Loss: 0.5799 Accuracy: 0.9641 Time: 15.57718  | Val Loss: 0.6647 Accuracy: 0.9388\n",
      "Epoch: 130/150 Train Loss: 0.5836 Accuracy: 0.9631 Time: 15.89729  | Val Loss: 0.6642 Accuracy: 0.9384\n",
      "Epoch: 131/150 Train Loss: 0.5804 Accuracy: 0.9646 Time: 15.89297  | Val Loss: 0.6660 Accuracy: 0.9388\n",
      "Epoch: 132/150 Train Loss: 0.5797 Accuracy: 0.9649 Time: 16.13880  | Val Loss: 0.6611 Accuracy: 0.9394\n",
      "Epoch: 133/150 Train Loss: 0.5780 Accuracy: 0.9653 Time: 16.13059  | Val Loss: 0.6633 Accuracy: 0.9384\n",
      "Epoch: 134/150 Train Loss: 0.5801 Accuracy: 0.9645 Time: 15.62207  | Val Loss: 0.6644 Accuracy: 0.9378\n",
      "Epoch: 135/150 Train Loss: 0.5763 Accuracy: 0.9662 Time: 15.45259  | Val Loss: 0.6638 Accuracy: 0.9391\n",
      "Epoch: 136/150 Train Loss: 0.5780 Accuracy: 0.9655 Time: 16.18041  | Val Loss: 0.6627 Accuracy: 0.9389\n",
      "Epoch: 137/150 Train Loss: 0.5741 Accuracy: 0.9672 Time: 16.44573  | Val Loss: 0.6620 Accuracy: 0.9389\n",
      "Epoch: 138/150 Train Loss: 0.5755 Accuracy: 0.9672 Time: 16.29756  | Val Loss: 0.6624 Accuracy: 0.9382\n",
      "Epoch: 139/150 Train Loss: 0.5747 Accuracy: 0.9662 Time: 15.54513  | Val Loss: 0.6609 Accuracy: 0.9390\n",
      "Epoch: 140/150 Train Loss: 0.5744 Accuracy: 0.9671 Time: 15.83630  | Val Loss: 0.6625 Accuracy: 0.9401\n",
      "Epoch: 141/150 Train Loss: 0.5717 Accuracy: 0.9684 Time: 16.03483  | Val Loss: 0.6594 Accuracy: 0.9407\n",
      "Epoch: 142/150 Train Loss: 0.5714 Accuracy: 0.9685 Time: 16.04492  | Val Loss: 0.6583 Accuracy: 0.9409\n",
      "Epoch: 143/150 Train Loss: 0.5712 Accuracy: 0.9683 Time: 15.69882  | Val Loss: 0.6585 Accuracy: 0.9406\n",
      "Epoch: 144/150 Train Loss: 0.5716 Accuracy: 0.9677 Time: 16.17661  | Val Loss: 0.6583 Accuracy: 0.9402\n",
      "Epoch: 145/150 Train Loss: 0.5690 Accuracy: 0.9693 Time: 15.93553  | Val Loss: 0.6553 Accuracy: 0.9431\n",
      "Epoch: 146/150 Train Loss: 0.5697 Accuracy: 0.9696 Time: 15.63444  | Val Loss: 0.6557 Accuracy: 0.9413\n",
      "Epoch: 147/150 Train Loss: 0.5713 Accuracy: 0.9682 Time: 15.57896  | Val Loss: 0.6564 Accuracy: 0.9424\n",
      "Epoch: 148/150 Train Loss: 0.5726 Accuracy: 0.9674 Time: 15.60046  | Val Loss: 0.6549 Accuracy: 0.9424\n",
      "Epoch: 149/150 Train Loss: 0.5700 Accuracy: 0.9687 Time: 16.16355  | Val Loss: 0.6554 Accuracy: 0.9416\n",
      "Epoch: 150/150 Train Loss: 0.5704 Accuracy: 0.9688 Time: 15.71516  | Val Loss: 0.6551 Accuracy: 0.9416\n"
     ]
    }
   ],
   "source": [
    "from hdd.train.warmup_scheduler import GradualWarmupScheduler\n",
    "from hdd.models.transformer.cswin_transformer import CSWinTransformer\n",
    "from hdd.train.classification_utils import naive_train_classification_model\n",
    "from hdd.models.nn_utils import count_trainable_parameter\n",
    "\n",
    "\n",
    "net = CSWinTransformer(\n",
    "    img_size=32,\n",
    "    in_chans=3,\n",
    "    num_classes=10,\n",
    "    embed_dim=48,\n",
    "    depth=[1, 2, 8, 1],\n",
    "    split_size=[1, 2, 2, 2],\n",
    "    num_heads=[2, 4, 8, 8],\n",
    "    mlp_ratio=4.0,\n",
    "    qkv_bias=True,\n",
    "    qk_scale=None,\n",
    "    drop_rate=0.0,\n",
    "    attn_drop_rate=0.0,\n",
    "    drop_path_rate=0.0,\n",
    "    norm_layer=nn.LayerNorm,\n",
    "    use_chk=False,\n",
    ").to(DEVICE)\n",
    "print(f\"#Parameter: {count_trainable_parameter(net)}\")\n",
    "criteria = nn.CrossEntropyLoss(label_smoothing=0.1)\n",
    "optimizer = torch.optim.Adam(\n",
    "    net.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=1e-5\n",
    ")\n",
    "\n",
    "base_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
    "    optimizer, max_epochs, eta_min=1e-5\n",
    ")\n",
    "scheduler = GradualWarmupScheduler(\n",
    "    optimizer,\n",
    "    multiplier=1.0,\n",
    "    total_epoch=10,\n",
    "    after_scheduler=base_scheduler,\n",
    ")\n",
    "random_setting = naive_train_classification_model(\n",
    "    net,\n",
    "    criteria,\n",
    "    max_epochs,\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    DEVICE,\n",
    "    optimizer,\n",
    "    scheduler,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b9d5da3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "6e13e8f8",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38d9724d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-cu124",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
